# sort_for_overlap.py
from matplotlib import pyplot as plt

from visual_image import modern_dbscan_overlap_data
from sort_API import get_overlap_rate
from config import DISPLAY_OVERLAP
import re


def extract_chinese_characters(word):
    """从字符串中提取所有汉字字符"""
    chinese_word = ''
    for char in word:
        if re.match(r'^[\u4e00-\u9fff]$', char):
            chinese_word += char
    return chinese_word


def is_likely_pinyin_line(line_data, pinyin_threshold):
    """
    判断一行是否可能是拼音行

    Args:
        line_data: 一行字符数据
        pinyin_threshold: 拼音字符比例阈值

    Returns:
        bool: 是否是拼音行
    """
    if not line_data:
        return False

    pinyin_count = 0
    total_count = len(line_data)

    for item in line_data:
        word = item['word']
        chinese_only = extract_chinese_characters(word)
        # 如果提取后为空，说明是拼音或非汉字字符
        if not chinese_only:
            pinyin_count += 1

    pinyin_ratio = pinyin_count / total_count
    return pinyin_ratio >= pinyin_threshold


def filter_pinyin_lines(sorted_data, pinyin_threshold):
    """
    过滤掉拼音行，只保留汉字行

    Args:
        sorted_data: 排序后的数据
        pinyin_threshold: 拼音字符比例阈值

    Returns:
        list: 过滤后的数据
    """
    filtered_data = []
    removed_lines_count = 0

    for column_idx, column_lines in enumerate(sorted_data):
        filtered_column = []

        for line_idx, line_data in enumerate(column_lines):
            if not is_likely_pinyin_line(line_data, pinyin_threshold):
                filtered_column.append(line_data)
            else:
                # print(f"移除第 {column_idx + 1} 栏第 {line_idx + 1} 行: 拼音行")
                removed_lines_count += 1

        filtered_data.append(filtered_column)

    # print(f"共移除 {removed_lines_count} 个拼音行")
    return filtered_data


def cleanup_chinese_characters(sorted_data):
    """
    清理每个字符，只保留汉字部分
    """
    cleaned_data = []

    for column_lines in sorted_data:
        cleaned_column = []

        for line_data in column_lines:
            cleaned_line = []
            for item in line_data:
                original_word = item['word']
                chinese_only = extract_chinese_characters(original_word)

                if chinese_only:  # 只保留包含汉字的字符
                    cleaned_item = item.copy()
                    cleaned_item['word'] = chinese_only
                    cleaned_line.append(cleaned_item)

            if cleaned_line:  # 只保留非空的行
                cleaned_column.append(cleaned_line)

        cleaned_data.append(cleaned_column)

    return cleaned_data


def overlap_sort(sorted_data_dbscan):
    sorted_data = []

    # 遍历每个栏（保持DBSCAN的分栏结果）
    for column_idx, column_data in enumerate(sorted_data_dbscan):
        if not column_data:
            sorted_data.append([])  # 空栏
            continue

        # print(f"处理第 {column_idx + 1} 栏，共 {len(column_data)} 个字符")

        # 确保栏内数据按y坐标排序（从上到下）
        column_data.sort(key=lambda x: x['y_start'])

        # 在当前栏内进行行分组（基于y轴重叠率）
        column_lines = []
        current_line = [column_data[0]]  # 当前行包含第一个元素

        for idx in range(1, len(column_data)):
            current_elem = column_data[idx]
            last_elem_in_line = current_line[-1]  # 与当前行最后一个元素比较

            # 计算y轴重叠率
            start1, end1 = current_elem['y_start'], current_elem['y_end']
            start2, end2 = last_elem_in_line['y_start'], last_elem_in_line['y_end']
            overlap = get_overlap_rate(start1, end1, start2, end2)

            if overlap >= 0.85:  # 重叠率阈值
                current_line.append(current_elem)
            else:
                # 对当前行按x坐标排序（从左到右）
                current_line.sort(key=lambda x: x['x_start'])
                column_lines.append(current_line)
                current_line = [current_elem]  # 开始新的一行

        # 添加最后一行
        if current_line:
            current_line.sort(key=lambda x: x['x_start'])
            column_lines.append(current_line)

        # 按y坐标对当前栏的所有行进行排序（从上到下）
        column_lines.sort(key=lambda line: line[0]['y_start'])

        # 添加到结果中（保持分栏结构）
        sorted_data.append(column_lines)

        # print(f"  第 {column_idx + 1} 栏分成 {len(column_lines)} 行")

    # 过滤拼音行
    # print("\n开始过滤拼音行...")
    filtered_data = filter_pinyin_lines(sorted_data, pinyin_threshold=0.6)

    # 清理汉字字符（移除非汉字部分）
    # print("清理汉字字符...")
    cleaned_data = cleanup_chinese_characters(filtered_data)

    # 移除空行和空栏
    final_data = []
    for column_lines in cleaned_data:
        non_empty_lines = [line for line in column_lines if line]
        if non_empty_lines:
            final_data.append(non_empty_lines)

    # 输出最终统计
    total_lines = sum(len(column) for column in final_data)
    total_chars = sum(len(char) for column in final_data for line in column for char in line)
    # print(f"\n最终结果: {len(final_data)} 栏, {total_lines} 行, {total_chars} 个汉字字符")

    return final_data
